#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-10-04

"""
【节点名称】：
    SpireFormatReaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    resize (list): 图像重采样分辨率
    fps (int): 输出帧率
【备注】：
    无
"""
import os
import threading
import time
import json
import cv2
import uuid
import platform
import argparse
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, def_msg, QoS, BaseNode, get_extra_args


def load_spire_dir(spire_dir, with_relative_path=False):
    """
    :param spire_dir: path to spire dataset root dir
    :param with_relative_path: output relative path or not
    :return: List - img_list, List - ann_list
    """
    img_list = []
    ann_list = []
    rel_img_list = []
    rel_ann_list = []
    for sub_dir in os.listdir(spire_dir):
        sub_fn = os.path.join(spire_dir, sub_dir)
        if os.path.isdir(sub_fn):
            if sub_dir == 'scaled_images':
                img_list.extend([os.path.join(sub_fn, img_nm) for img_nm in os.listdir(sub_fn) if
                                 os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                ann_list.extend([os.path.join(spire_dir, 'annotations', img_nm + '.json') for img_nm in
                                 os.listdir(sub_fn) if
                                 os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                if with_relative_path:
                    rel_img_list.extend([os.path.join(sub_dir, img_nm) for img_nm in os.listdir(sub_fn) if
                                        os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                    rel_ann_list.extend([os.path.join('annotations', img_nm + '.json') for img_nm in
                                        os.listdir(sub_fn) if
                                        os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
            else:
                for sub_sub_dir in os.listdir(sub_fn):
                    sub_sub_fn = os.path.join(sub_fn, sub_sub_dir)
                    if os.path.isdir(sub_sub_fn):
                        if sub_sub_dir == 'scaled_images':
                            img_list.extend([os.path.join(sub_sub_fn, img_nm) for img_nm in os.listdir(sub_sub_fn) if
                                             os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                            ann_list.extend([os.path.join(sub_fn, 'annotations', img_nm + '.json') for img_nm in
                                             os.listdir(sub_sub_fn) if
                                             os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                            if with_relative_path:
                                rel_img_list.extend([os.path.join(sub_dir, sub_sub_dir, img_nm)
                                                     for img_nm in os.listdir(sub_sub_fn)
                                                     if os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                                rel_ann_list.extend([os.path.join(sub_dir, 'annotations', img_nm + '.json')
                                                     for img_nm in os.listdir(sub_sub_fn)
                                                     if os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
    if with_relative_path:
        img_list.sort()
        ann_list.sort()
        rel_img_list.sort()
        rel_ann_list.sort()
        return img_list, ann_list, rel_img_list, rel_ann_list
    else:
        img_list.sort()
        ann_list.sort()
        return img_list, ann_list


class SpireFormatReaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True,
            **kwargs
        )
        self.spire_root = self.get_param("spire_root", "D:/dataset/GT")
        self.dataset_name = self.get_param("dataset_name", "coco_detection")
        self.categories = self.get_param("categories", ["boat"])
        self.next_auto = self.get_param("next_auto", 0)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.sms_shutdown_emit = self.get_param("sms_shutdown_emit", True)
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.queue_pool.append(self.img_i_queue)
        self.img_i = 0
        self.frame_id = 0

        evaluate_url = '/SpireView/{}/TaskResults'.format(self.client_id)
        self._evaluate_reader = Subscriber(
            evaluate_url, 'spirecv_msgs::EvaluationResult', self.evaluate_callback,
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            '/' + job_name + '/gt/results', 'spirecv_msgs::2DTargets',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._shutdown_writer = Publisher(
            '/' + job_name + '/shutdown', 'std_msgs::Boolean',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._evaluate_writer = Publisher(
            '/SpireView/EvaluationResult', 'spirecv_msgs::EvaluationResult',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self.img_list, self.ann_list, self.rel_img_list, self.rel_ann_list = load_spire_dir(
            self.spire_root, with_relative_path=True)

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._next_reader.kill()
        self._result_writer.kill()
        self._evaluate_reader.kill()
        self._evaluate_writer.kill()
        self._shutdown_writer.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.img_list):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1
            else:
                print('Finish! {} / {}'.format(self.img_i, len(self.img_list)))

    def evaluate_callback(self, msg):
        print(msg)
        self._evaluate_writer.publish(msg)
        if self.sms_shutdown_emit:
            msg = def_msg('std_msgs::Boolean')
            msg['data'] = True
            self._shutdown_writer.publish(msg)

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img = cv2.imread(self.img_list[img_i])
            rel_path = self.rel_img_list[img_i]
            parts = os.path.normpath(rel_path).split(os.sep)

            with open(self.ann_list[img_i], 'r') as f:
                gt = json.load(f)

            assert parts[-1] == gt['file_name']

            tgt_msg = def_msg('spirecv_msgs::2DTargets')
            tgt_msg['file_name'] = gt['file_name']
            tgt_msg['height'] = gt['height']
            tgt_msg['width'] = gt['width']
            tgt_msg['dataset'] = self.dataset_name
            tgt_msg['client_id'] = self.client_id
            tgt_msg['img_id'] = img_i
            tgt_msg['img_total'] = len(self.img_list)
            tgt_msg['targets'] = []
            for anno in gt['annos']:
                if anno['category_name'] in self.categories:
                    tgt_msg['targets'].append(anno)
            self._result_writer.publish(tgt_msg)
            # print(tgt_msg)

            if self.b_use_shm:
                msg = self._image_writer.cvimg2sms_mem(img)
            else:
                msg = cvimg2sms(img)

            msg['img_id'] = img_i
            msg['img_total'] = len(self.img_list)
            msg['file_name'] = os.path.basename(self.img_list[img_i])
            msg['client_id'] = self.client_id
            if len(parts) == 3:
                msg['video_name'] = parts[0]

            msg['spirecv_msgs::2DTargets'] = tgt_msg
            self._image_writer.publish(msg)

            if self.next_auto:
                time.sleep(1)
                self.launch_next()

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = SpireFormatReaderNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.launch_next()
    node.join()
